#include "logit_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;

static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    LogitCustomTilingData tiling;
    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    auto coreNum = ascendcPlatform.GetCoreNum();
    
    // 获取输入数据信息
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t typeLength = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);
    
    // 修正：正确计算字节数
    uint32_t inputBytes = typeLength;  // 每个数据元素的字节数

    // 计算UB使用情况：输入队列(2个buffer) + 输出队列(2个buffer) + 4个临时buffer = 8个buffer
    uint32_t totalBufferNum = 2 * BUFFER_NUM + 4;  // inQueue + outQueue + 4 tmpBuf
    uint32_t tileDataNum = (ubSize / totalBufferNum) / inputBytes;
    
    // 32B对齐
    tileDataNum = (tileDataNum / BLOCK_SIZE) * BLOCK_SIZE / inputBytes;
    if (tileDataNum == 0) {
        tileDataNum = BLOCK_SIZE / inputBytes;  // 至少保证一个block的数据量
    }

    // 计算每个核心处理的数据量
    uint32_t totalDataNum = inputNum;
    uint32_t averageDataPerCore = (totalDataNum + coreNum - 1) / coreNum;  // 向上取整
    
    // 确保有足够的核心参与计算
    if (averageDataPerCore == 0) {
        coreNum = 1;
        averageDataPerCore = totalDataNum;
    }
    
    // 计算有多少个核心需要处理额外的数据
    uint32_t tailBlockNum = totalDataNum % coreNum;
    uint32_t smallCoreDataNum = totalDataNum / coreNum;
    uint32_t bigCoreDataNum = smallCoreDataNum + (tailBlockNum > 0 ? 1 : 0);
    
    // 如果有剩余数据，前tailBlockNum个核心多处理1个数据
    if (tailBlockNum == 0) {
        tailBlockNum = coreNum;  // 所有核心处理相同数量的数据
        bigCoreDataNum = smallCoreDataNum;
    }
    
    // 计算每个核心的tile数量
    uint32_t smallTileNum = (smallCoreDataNum + tileDataNum - 1) / tileDataNum;
    uint32_t bigTileNum = (bigCoreDataNum + tileDataNum - 1) / tileDataNum;
    
    // 计算尾部数据量
    uint32_t smallTailDataNum = smallCoreDataNum % tileDataNum;
    if (smallTailDataNum == 0 && smallCoreDataNum > 0) {
        smallTailDataNum = tileDataNum;
    }
    
    uint32_t bigTailDataNum = bigCoreDataNum % tileDataNum;
    if (bigTailDataNum == 0 && bigCoreDataNum > 0) {
        bigTailDataNum = tileDataNum;
    }
    
    tiling.set_smallCoreDataNum(smallCoreDataNum);
    tiling.set_bigCoreDataNum(bigCoreDataNum);
    tiling.set_tileDataNum(tileDataNum);
    tiling.set_smallTailDataNum(smallTailDataNum);
    tiling.set_bigTailDataNum(bigTailDataNum);
    tiling.set_finalSmallTileNum(smallTileNum);
    tiling.set_finalBigTileNum(bigTileNum);
    tiling.set_tailBlockNum(tailBlockNum);
    
    context->SetBlockDim(coreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}

static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
}

namespace ops {
class LogitCustom : public OpDef {
public:
    explicit LogitCustom(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend310b")
            .AddConfig("ascend910b");
    }
};

OP_ADD(LogitCustom);
}